import sys
import os.path as osp
import os
sys.path.append( osp.join(os.getcwd(),"/model/extractors/")) 


import torch
from torch import nn
from .nets.sfd2 import ResSegNet, ResSegNetV2
from .nets.extractor import extract_resnet_return

def get_model(model_name, weight_path, use_stability=False):
    if model_name == 'ressegnet':
        model = ResSegNet(outdim=128, require_stability=use_stability).eval()
        model.load_state_dict(torch.load(weight_path)['model'], strict=True)
        extractor = extract_resnet_return
    if model_name == 'ressegnetv2':
        model = ResSegNetV2(outdim=128, require_stability=use_stability).eval()
        model.load_state_dict(torch.load(weight_path)['model'], strict=False)
        extractor = extract_resnet_return
    return model, extractor

class Features_SFD2(nn.Module):
    def __init__(self, *args):
        super().__init__()
        conf = {
            'output': 'feats-ressegnetv2-20220810-wapv2-sd2mfsf-uspg-0001-n3000-r1600',
            'model': {
                'name': 'ressegnetv2',
                'use_stability': True,
                'max_keypoints': 3000,
                'conf_th': 0.001,
                'multiscale': False,
                'scales': [1.0],
                'model_fn': "checkpoints/20220810_ressegnetv2_wapv2_ce_sd2mfsf_uspg.pth",
            },
            'preprocessing': {
                'grayscale': False,
                'resize_max': 1600,
            },
            'mask': False,
        }
        model, extractor = get_model(model_name=conf['model']['name'], weight_path=conf["model"]["model_fn"],
                                    use_stability=conf["model"]['use_stability'])
        
        model = model.cuda()
        self.model = model
        
        self.extractor = extractor
        self.conf = conf

    def forward(self, x):
        # getter = cv2.SIFT_create()
        pred = self.extractor(self.model, img=x,
                    topK=self.conf["model"]["max_keypoints"],
                    mask=None,
                    conf_th=self.conf["model"]["conf_th"],
                    scales=self.conf["model"]["scales"],
                    )
        pts = pred["keypoints"]
        desc = pred["descriptors"]
        return {"keypoints":pts.transpose(),"descriptors":desc.transpose()}